package com.edu.robot.core.service.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.embeddings.TextEmbeddingParam;
import com.alibaba.dashscope.embeddings.TextEmbeddingResult;
import com.alibaba.dashscope.embeddings.TextEmbeddingResultItem;
import com.alibaba.fastjson2.JSON;
import com.edu.robot.core.enums.IntentEnum;
import com.edu.robot.core.enums.NotifyMsgEnum;
import com.edu.robot.core.model.*;
import com.edu.robot.core.service.ModelManger;
import com.edu.robot.modules.robot.dto.ModelConfigDTO;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @program: KQAIA
 * @description:
 * @author: 无恙
 * @create: 2025-03-17 19:23
 **/

@Service
@Slf4j
public class ModelMangerImpl implements ModelManger {


	String regex = "intent: \"(.*?)\",\\s*confidence: (\\d+\\.?\\d*)";

	@Value("${model.key:sk-a4298b6daa5c4efb940d11518e960a65}")
	private String APIKEY;

	@Override
	@SneakyThrows(Exception.class)
	public Integer textEmbedding(String text,List<Float> vector) {

		if (text == null || text.isEmpty()) {
			log.error("textEmbedding vectorResponses is empty");
			return 0;
		}

		TextEmbeddingParam param = TextEmbeddingParam
				.builder()
				.model("text-embedding-v3")
				.apiKey(APIKEY)
				//设置向量维度
				.parameters(Collections.singletonMap("dimension", "512"))
				.texts(Collections.singleton(text)).build();
		TextEmbedding textEmbedding = new TextEmbedding();
		TextEmbeddingResult result = textEmbedding.call(param);
		result.getOutput().getEmbeddings().stream().map(TextEmbeddingResultItem::getEmbedding).findFirst().ifPresent(
				embeddings -> embeddings.stream().map(Double::floatValue).forEach(vector::add)
		);

		return result.getUsage().getTotalTokens();

	}

	@Override
	@SneakyThrows(Exception.class)
	public String intentRecognition(QueryContext queryContext) {

		Generation gen = new Generation();
		Message systemMsg = Message.builder()
				.role(Role.SYSTEM.getValue())
				.content(IntentEnum.INTENT_AGENT.getSysPrompt())
				.build();
		Message userMsg = Message.builder()
				.role(Role.USER.getValue())
				.content(queryContext.getUserInput())
				.build();
		GenerationParam param = GenerationParam.builder()
				// 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
				.apiKey(APIKEY)
				// 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
				.model("qwen-plus")
				.messages(Arrays.asList(systemMsg, userMsg))
				.resultFormat(GenerationParam.ResultFormat.TEXT)
				.build();
		String text = gen.call(param).getOutput().getText();
		Integer totalTokens = gen.call(param).getUsage().getTotalTokens();
		log.info("intentRecognition text:{},totalTokens:{}",text,totalTokens);

		Pattern pattern = Pattern.compile(regex);
		Matcher matcher = pattern.matcher(text);

		if (matcher.find()) {
			String intent = matcher.group(1);    // 提取第1个分组（intent值）
			String confidence = matcher.group(2); // 提取第2个分组（confidence值）

			//TODO 多agent校验
			return intent;
		} else {
			log.error("intentRecognition no match");
		}
		return null;
	}

	@Override
	@SneakyThrows
	public String modelInvoke(SessionContext context) {

		Generation gen = new Generation();
		Message userMsg = Message.builder()
				.role(Role.USER.getValue())
				.content(context.getQuestion())
				.build();
		GenerationParam param = GenerationParam.builder()
				// 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
				.apiKey(APIKEY)
				// 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
				.model("qwen-max")
				.prompt(context.getPrompt())
				.messages(Collections.singletonList(userMsg))
				.resultFormat(GenerationParam.ResultFormat.TEXT)
				.build();

		String text = gen.call(param).getOutput().getText();
		Integer totalTokens = gen.call(param).getUsage().getTotalTokens();

		context.setFinalAnswer(text);
		context.setToken(totalTokens);
		log.info("intentRecognition text:{},totalTokens:{}",text,totalTokens);
		return text;
	}

	@Override
	@SneakyThrows
	public String modelInvokeStream(SessionContext context) {
		ModelConfigDTO modelConfigDTO = context.getModelConfigDTO();
		Message userMsg = Message.builder()
				.role(Role.USER.getValue())
				.content(context.getQuestion())
				.build();

		GenerationParam param = GenerationParam.builder()
				// 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
				.apiKey(APIKEY)
				// 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
				.model(modelConfigDTO.getModelName())
				.prompt(context.getPrompt())
				.messages(Collections.singletonList(userMsg))
				.resultFormat(GenerationParam.ResultFormat.MESSAGE)
				.build();
		Optional.ofNullable(modelConfigDTO.getTopP()).ifPresent(param::setTopP);
		Optional.ofNullable(modelConfigDTO.getTopK()).ifPresent(param::setTopK);
		Optional.ofNullable(modelConfigDTO.getTemperature()).ifPresent(param::setTemperature);
		Optional.ofNullable(modelConfigDTO.getStopWords()).ifPresent(s->param.setStopStrings(Collections.singletonList(s)));


		Generation gen = new Generation();

		APICallback callback = context.getCallback();


		//计费
		AtomicInteger atomicInteger = new AtomicInteger(0);
		Semaphore sem = new Semaphore(0);
		gen.streamCall(param, new ResultCallback<GenerationResult>() {
			@Override
			public void onEvent(GenerationResult generationResult) {
				String content = generationResult.getOutput().getChoices().get(0).getMessage().getContent();
				context.setFinalAnswer(content);
				atomicInteger.addAndGet(generationResult.getUsage().getTotalTokens());
				callback.onEvent(NotifyMsgEnum.NORMAL,content);

			}

			@Override
			public void onComplete() {

				sem.release();
			}

			@Override
			public void onError(Exception e) {
				log.error("onError:{}",e.getMessage());

				callback.onError(e);

				sem.release();
			}
		});

		sem.acquire();
		if(context.getToken() == 0){
			context.setToken(atomicInteger.get());
		}else {
			context.setToken(context.getToken() + atomicInteger.get());
		}

		log.info("modelInvokeStream text:{},totalTokens:{}",context.getFinalAnswer().toString(),atomicInteger.get());

		callback.close(context.getFinalAnswer());
		return context.getFinalAnswer();
	}

	@Override
	@SneakyThrows(Exception.class)
	public String titleRecognition(String query) {

		Generation gen = new Generation();
		Message systemMsg = Message.builder()
				.role(Role.SYSTEM.getValue())
				.content(IntentEnum.INTENT_TITLE.getSysPrompt())
				.build();
		Message userMsg = Message.builder()
				.role(Role.USER.getValue())
				.content(query)
				.build();
		GenerationParam param = GenerationParam.builder()
				// 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
				.apiKey(APIKEY)
				// 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
				.model("qwen-plus")
				.messages(Arrays.asList(systemMsg, userMsg))
				.resultFormat(GenerationParam.ResultFormat.TEXT)
				.build();
		return gen.call(param).getOutput().getText();
	}


}
